import torch
from tqdm import tqdm




def score(h, r, t):
    s = torch.sum((h * r) * t, -1).flatten().item()
    return s


def quit_sort2(arr, left, right):
    """快速排序算法
        思想：快速排序是效率最高的排序算法之一,
        1.先以一个基准开始（通常为第一个数）
        2.从右边开始找到第一个比基准小的数
        3.从左边开始找到第一个比基准大的数
        4.将比基准小的数和比基准大的数调换位置
        5.继续从右往左以及从左往右调到对应的数，直到二边重合结束
        6.将基准的数与重合的数调换位置，这是基准左边都是小的数，右边都是大的数
        7.重复1,2,3,4对左右两边的数进行排序
    """
    # 列表长度大于1
    if len(arr) <= 1 or left > right:
        return
    # 定义基准
    temp = arr[left]
    # 定义从左边开始的数(基准的下一个数)
    i = left
    # 定义从右边开始的数（要判断的末尾）
    j = right
    # 一直循环到从左边开始的碰到从右边开始的
    while i < j:
        # 从右边开始的一直循环到第一个比基准小的数且没有碰头为止
        while arr[j][0] <= temp[0] and i < j:
            j -= 1
        # 从左边开始的一直循环到第一个比基准大的数且没有碰头为止
        while arr[i][0] >= temp[0] and i < j:
            i += 1
        # 当在两边找到符合条件且没有碰头的数据时，将数据调换位置
        if i < j:
            mTemp = arr[j]
            arr[j] = arr[i]
            arr[i] = mTemp
        # 继续往下循环找到所有的符合条件的并调换位置

    # 到这里当前基准循环完毕，将当前基准和结束循环的位置数调换，使得基准左边都是小于基准的数，右边都是大于基准的数
    arr[left] = arr[i]
    arr[i] = temp
    # 继续将基准两边的数进行以上操作
    quit_sort2(arr, left, j - 1)
    quit_sort2(arr, j + 1, right)


f = open("benchmarks/MY/患者电子病历数据集/dataset.txt", "r")
lines = f.readlines()
EMRs = []
for line in lines[1:]:
    EMRs.append(line.split())
f.close()

f = open("benchmarks/MY/医学知识/uncertainty.txt")
lines = f.readlines()
weights = []
for line in lines[1:]:
    weights.append(line.split())
f.close()

# print(len(weights))
# print(len(EMRs))

ckpt = torch.load('./checkpoint/distmult.ckpt')
ents = ckpt['ent_embeddings.weight']  # [0:27]疾病， [26:]症状
rels = ckpt['rel_embeddings.weight']
# print(len(ents))

right_num10 = 0
right_num3 = 0
right_num1 = 0
count = 0
bar = tqdm(EMRs)
for EMR in bar:
    # print(EMR)
    hs = []
    scores = []
    for e in EMR[:-1]:
        hs.append([ents[int(e)], int(e)])
    # print(hs)

    # 遍历每一种疾病，给这个疾病打分
    for i in range(0, 27):
        this_score = 0.0
        # 遍历症状集中每个症状，每个症状与此疾病的分数加在一起
        for h in hs:
            this_score += score(h[0], rels[0], ents[i])
        scores.append([this_score, i])
    quit_sort2(scores, 0, len(scores) - 1)
    # print(scores)
    for k in range(0, 10):
        if int(EMR[-1]) == scores[k][1] and k < 1:
            right_num1 += 1
        if int(EMR[-1]) == scores[k][1] and k < 3:
            right_num3 += 1
        if int(EMR[-1]) == scores[k][1] and k < 10:
            right_num10 += 1
    count += 1
    bar.set_description("前十：%f | 前十：%f | 第一：%f " % (right_num10 / count, right_num3 / count, right_num1 / count))
    # print("下一轮")
